
# 此代码用于生成导丝分割数据集

import simulation
import pygame
import os
import math
import json
from pygame.key import *
from pygame.locals import *
from pygame.color import *
import cv2

import numpy as np
import io

import pymunk
import pymunk.constraints
import pymunk.pygame_util

def num_compare(name:str):
    return int(name.split(".")[0])

def get_biggist_id(path:str) -> int:
    f_names = os.listdir(path)
    if len(f_names) == 0:
        return -1
    f_names.sort(key=num_compare)
    return int(f_names[-1].split(".")[0])

RADIUS = 5
DIR_ANGEL = 45
ID = 0
DIR = "./datas/test/"
ORIGINAL_DIR  = os.path.join(DIR, "images")
MASK_DIR      = os.path.join(DIR, "label")
MAP_JSON_PATH = os.path.join(DIR, "envmsgs")
G_IMAGE_PATH    = os.path.join(DIR, "g_image")
G_LABEL_PATH    = os.path.join(DIR, "g_label")


os.makedirs(ORIGINAL_DIR, exist_ok=True)
os.makedirs(MASK_DIR, exist_ok=True)
os.makedirs(MAP_JSON_PATH, exist_ok=True)
os.makedirs(G_IMAGE_PATH, exist_ok=True)
os.makedirs(G_LABEL_PATH, exist_ok=True)



G_LABEL_ID      = get_biggist_id(G_LABEL_PATH) + 1


G_RENDER_WIDTH  = 4 
G_RENDER_RGB  = (43, 43, 43) 

ORIGINAL_IMAGES = os.listdir(ORIGINAL_DIR)
MASK_IMAGES     = os.listdir(MASK_DIR)


ORIGINAL_IMAGES.sort(key=num_compare)
MASK_IMAGES.sort(key=num_compare)

ORIGINAL_IMAGES = [os.path.join(ORIGINAL_DIR, f_name) for f_name in ORIGINAL_IMAGES]
MASK_IMAGES = [os.path.join(MASK_DIR, f_name) for f_name in MASK_IMAGES]

_flag = 0
POINTS  = [None for _ in range(3)]
PTCOLORS  = ["blue", "red", "pink"]


def save_data(data, path):
    with open(path, "w") as f:
        json.dump(data, f)

def load_data(path):
    with open(path, "r") as f:
        data = json.load(f)
    return data




def surf_to_ndarray(surf:pygame.Surface, alpha=False):
    """pygame image to cv2 form image"""
    frame_np = np.array(pygame.surfarray.pixels3d(surf))
    if alpha:
        alpha_np = np.array(pygame.surfarray.pixels_alpha(surf))[:,:,np.newaxis]
        # alpha_np = np.array(pygame.surfarray.pixels_alpha(surf))
        frame_np = np.append(frame_np, alpha_np, axis=2)
        # frame_np = np.stack((frame_np, alpha_np), axis=2)
    frame_np = np.transpose(frame_np, (1, 0, 2))
    return frame_np

def ndarray_to_surf(ndarray:np.ndarray):
    conv_img = pygame.surfarray.make_surface(ndarray[:,:,:3])
    # conv_img.set_alpha(pygame.surfarray.)
    conv_img = pygame.transform.rotate(conv_img, 90)
    return pygame.transform.flip(conv_img,False,True)

def angle_between_points(point1, point2):
    x_diff = point2[0] - point1[0]
    y_diff = point2[1] - point1[1]
    return math.degrees(math.atan2(y_diff, x_diff))


class App:
    def __init__(self):
        pygame.init()
        self.env = simulation.GuideWireEngine()
        self.bg   = pygame.image.load(ORIGINAL_IMAGES[ID])
        self.mask = pygame.image.load(MASK_IMAGES[ID])
        self.bg   = pygame.transform.scale(self.bg, (self.mask.get_size()[0] *2, self.mask.get_size()[1] * 2))
        self.mask = pygame.transform.scale(self.mask, (self.mask.get_size()[0] *2, self.mask.get_size()[1] * 2))
        
         

        self.env.guidwire_fric = 0.1
        self.screen = pygame.Surface(self.mask.get_size(), pygame.SRCALPHA)
        self.display = pygame.display.set_mode(self.mask.get_size())
        self.draw_options = pymunk.pygame_util.DrawOptions(self.display)
        # pymunk.pygame_util.positive_y_is_up = False
        # with open("real.json") as json_file:
        #     self.wall_poslist = json.load(json_file)	
        self.insert_limit = 100
        # 载入血管壁
        self.env.draw_walls_by_mask(self.mask, friction=0.1)
        # 载入导丝
        # self.env.create_goal((310, 485))
        # self.env.create_ini_guidewire(radius=RADIUS, ini_len=1, start_position=(69, 113), direct_angle=DIR_ANGEL)


    def events_handler(self):
        global POINTS, _flag, DIR_ANGEL
        for event in pygame.event.get():
            if event.type == QUIT:
                exit()
            elif event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    exit()
                if event.key == K_d:
                    if len(self.env.balls) <= self.insert_limit and _flag > 1:
                        self.env.insert_a_ball(radius=RADIUS, direct_angle=DIR_ANGEL)
                        if len(self.env.balls) >= 3:
                            self.save_g_data()
                if event.key == K_a:
                    self.env.pull_back_a_ball(direct_angle=DIR_ANGEL)
                    if len(self.env.balls) >= 3:
                        self.save_g_data()
                if event.key == K_w:
                    self.env.bend_guidewire(d_angle=10)
                if event.key == K_s:
                    self.env.bend_guidewire(d_angle=-10)
                
                if event.key == K_p:
                    print(surf_to_ndarray(self.screen).shape)
                    
                    cv2.imwrite(f"out.jpg", surf_to_ndarray(self.screen))
                if event.key == K_RIGHT:
                    self.load_next()
                    pygame.display.set_caption(f"File ID:{ID}")
                if event.key == K_LEFT:
                    self.load_prev()
                    pygame.display.set_caption(f"File ID:{ID}")
            elif event.type == MOUSEMOTION:
                pygame.display.set_caption(f"File ID:{ID}, Mouse At: {event.pos}")
            
            elif event.type == pygame.MOUSEBUTTONDOWN:
                ...
                # if _flag < 3:
                #     POINTS[_flag] = event.pos
                
                # if _flag == 0:
                #     ...
                #     # self.env.balls[0].body.position = event.pos
                
                # elif _flag == 1:
                #     DIR_ANGEL = angle_between_points(POINTS[0], POINTS[1])
                #     self.env.create_ini_guidewire(RADIUS, POINTS[0], 1, DIR_ANGEL)
                #     print(DIR_ANGEL)
                
                # elif _flag == 2:
                #     ...
                #     # self.env.create_goal(event.pos)


                # _flag += 1
                # if _flag > 3:
                #     save_data(POINTS, os.path.join(JSON_PATH, f"{ID}.json"))
                #     _flag = 0
                #     POINTS  = [None for _ in range(3)]
                #     self.load_next()

    def save_g_data(self):
        '''  '''
        global G_LABEL_ID

        surf_g_label = pygame.Surface(self.display.get_size())
        surf_g_label.fill("black")
        self.draw_guidewire(surf_g_label, "white", 5)
        pygame.image.save(self.display, os.path.join(G_IMAGE_PATH, f"{G_LABEL_ID}.png"))
        pygame.image.save(surf_g_label, os.path.join(G_LABEL_PATH, f"{G_LABEL_ID}.png"))
        G_LABEL_ID += 1


    def draw_guidewire(self, surface:pygame.Surface, color=None, width=None,):
        """ 绘制导丝 """
        g_color = G_RENDER_RGB if color is None else color
        g_width = G_RENDER_WIDTH if width is None else width

        points = []
        # 获取导丝的位姿信息并绘制导丝
        for i in self.env.balls:
            points.append(i.body.position.int_tuple)
        if len(points) > 1:
            pygame.draw.lines(surface, g_color, False, points, g_width)
            pygame.draw.circle(surface, g_color, points[-1], g_width//2)


    def render(self):

        self.display.blit(self.bg, (0,0))
        self.screen.fill((255,255,255, 0))
        # self.display.blit(self.mask, (0,0))
        self.draw_guidewire(self.screen)
        # points = []
        # 获取导丝的位姿信息并绘制导丝
        # for i in self.env.balls:
        #     points.append(i.body.position.int_tuple)
        # if len(points) > 1:
        #     pygame.draw.lines(self.screen, G_RENDER_RGB, False, points, G_RENDER_WIDTH)
        #     pygame.draw.circle(self.screen, G_RENDER_RGB, points[-1], G_RENDER_WIDTH)
            # 绘制人工智能识别的结果
            # w, h = 40, 40
            # rect = (points[-1][0] - w/2, points[-1][1] - h/2, w, h)
            # pygame.draw.rect(self.screen, (0, 255, 0), rect, 5)

        # 绘制肿瘤
        # for i in self.env.goals:		
        #     p = self.env.goals[-1].body.position.int_tuple
        #     pygame.draw.circle(self.display, (70, 27, 44), p, 15)
        # if len(self.env.goals) > 0:    
        #     rect = (p[0] - w/2, p[1] - h/2, w, h)
            # pygame.draw.rect(self.display, (255, 255, 0), rect, 5) 


        # 刷新屏幕
        # self.env.space.debug_draw(self.draw_options)  # 绘制新图案
        # self.display.blit(self.screen, (0,0))
        blured_arr = cv2.GaussianBlur(surf_to_ndarray(self.screen, True), (5, 5), 0) 
        img_encode_bytes = cv2.imencode(".png", blured_arr)[1].tobytes()
        picture_stream  = io.BytesIO(img_encode_bytes)

        # surf_encode = pygame.image.frombytes(img_encode, (blured.shape[1], blured.shape[0]), ".png")
        surf_blured = pygame.image.load(picture_stream, ".png")
        # self.screen = ndarray_to_surf(blured) 

        self.display.blit(self.bg, (0,0))
        # self.display.blit(self.screen, (0,0))
        self.display.blit(surf_blured, (0,0))
        # self.env.space.debug_draw(self.draw_options)  # 绘制新图案
        # 绘制鼠标点
        # for k, i in enumerate(POINTS):
        #     if i is not None:
        #         pygame.draw.circle(self.display, PTCOLORS[k], i, 3)

        pygame.display.flip()

	
    def debug_render(self):
        self.screen.blit(self.mask, (0,0))          # 清空图案
        self.env.space.debug_draw(self.draw_options)  # 绘制新图案
        self.display.blit(self.screen, (0,0))
        pygame.display.flip()

    def load_next(self):
        global ID, POINTS, _flag, DIR_ANGEL
        _flag = 0
        ID += 1
        POINTS  = [None for _ in range(3)]
        self.env.clear_all()
        self.bg   = pygame.image.load(ORIGINAL_IMAGES[ID])
        self.mask = pygame.image.load(MASK_IMAGES[ID])
        # self.screen.blit(self.mask, (0,0))
        self.bg   = pygame.transform.scale(self.bg, (self.mask.get_size()[0] *2, self.mask.get_size()[1] * 2))
        self.mask = pygame.transform.scale(self.mask, (self.mask.get_size()[0] *2, self.mask.get_size()[1] * 2))
        self.env.draw_walls_by_mask(self.mask, friction=0.1)

        if os.path.exists(os.path.join(MAP_JSON_PATH, f"{ID}.json")):
            POINTS = load_data(os.path.join(MAP_JSON_PATH, f"{ID}.json"))
            DIR_ANGEL = angle_between_points(POINTS[0], POINTS[1])
            self.env.create_ini_guidewire(RADIUS, POINTS[0], 1, DIR_ANGEL)
            _flag = 4
            print(DIR_ANGEL)

    def load_prev(self):
        global ID, POINTS, _flag, DIR_ANGEL
        _flag = 0
        ID = ID - 1 if not ID == 0 else ID
        self.env.clear_all()
        self.bg   = pygame.image.load(ORIGINAL_IMAGES[ID])
        self.mask = pygame.image.load(MASK_IMAGES[ID])
        # self.screen.blit(self.mask, (0,0))
        self.bg   = pygame.transform.scale(self.bg, (self.mask.get_size()[0] *2, self.mask.get_size()[1] * 2))
        self.mask = pygame.transform.scale(self.mask, (self.mask.get_size()[0] *2, self.mask.get_size()[1] * 2))
        self.env.draw_walls_by_mask(self.mask, friction=0.1)

        if os.path.exists(os.path.join(MAP_JSON_PATH, f"{ID}.json")):
            POINTS = load_data(os.path.join(MAP_JSON_PATH, f"{ID}.json"))
            DIR_ANGEL = angle_between_points(POINTS[0], POINTS[1])
            self.env.create_ini_guidewire(RADIUS, POINTS[0], 1, DIR_ANGEL)
            print(DIR_ANGEL)
            _flag = 4




def main():
    global ID, POINTS, _flag, DIR_ANGEL
    app = App()
    if os.path.exists(os.path.join(MAP_JSON_PATH, f"{ID}.json")):
        POINTS = load_data(os.path.join(MAP_JSON_PATH, f"{ID}.json"))
        DIR_ANGEL = angle_between_points(POINTS[0], POINTS[1])
        app.env.create_ini_guidewire(RADIUS, POINTS[0], 1, DIR_ANGEL)
        print(DIR_ANGEL)
        _flag = 4
    
    FPS = 60
    exact = 1000 / 60
    clock = pygame.time.Clock()
    while True:
        # 计算下一帧
        for x in range(int(exact)):
            app.env.space.step(1 / FPS / exact)
        app.render()
        # game.debug_render()
        app.events_handler()
        clock.tick(FPS)
		



if __name__ == '__main__':
	main()